from lab_4_prompt_entity_extract import get_entity_chain
from lab_5_create_neo4j_fulltxt_index import gen_full_txt_qry
from lab_2_conf_manage_neo4j import get_neo4j_graph


# Fulltext index query
def structured_retriever(question: str) -> str:
    result = ""
    entities = get_entity_chain().invoke({"question": question})
    print(entities)
    graph = get_neo4j_graph()
    for entity in entities.names:
        response = graph.query(
            """CALL db.index.fulltext.queryNodes('entity', $query, {limit:2})
            YIELD node,score
            CALL {
              MATCH (node)-[r:!MENTIONS]->(neighbor)
              RETURN node.id + ' - ' + type(r) + ' -> ' + neighbor.id AS output
              UNION
              MATCH (node)<-[r:!MENTIONS]-(neighbor)
              RETURN neighbor.id + ' - ' + type(r) + ' -> ' +  node.id AS output
            }
            RETURN output LIMIT 50
            """,
            {"query": gen_full_txt_qry(entity)},
        )
        result += "\n".join([el['output'] for el in response])
    return result


if __name__ == "__main__":
    print(structured_retriever("Who is Elizabeth I?"))
